from pathlib import Path
from functools import partial
import json

from tqdm import tqdm
from perfetto.trace_processor import TraceProcessor, TraceProcessorConfig


def get_utid_by_thread_name(thread_names: list[str] | str):
    if isinstance(thread_names, str):
        thread_names = [thread_names]
    for thread_name in thread_names:
        ui_thread = tp.query(f"select utid from thread where name = '{thread_name}' and is_main_thread=1")
        if len(ui_thread):
            break
    assert len(ui_thread) == 1, f"found {len(ui_thread)} ui thread(s), expecting 1"
    return list(ui_thread)[0].utid


def fetchone(sql: str, tp: TraceProcessor):
    res = list(tp.query(sql))
    assert len(res) == 1
    return res[0]


def fetchall(sql: str, tp: TraceProcessor):
    res = list(tp.query(sql))
    return res


trace_proc_config = TraceProcessorConfig(bin_path="/home/user/Desktop/software/perfetto/trace_processor_shell")


stats = []

for trace_path in Path("/home/user/Desktop/trace/vmall_heavy/").glob("*"):
    print(trace_path)
    tp = TraceProcessor(trace=str(trace_path), config=trace_proc_config)
    fetchone = partial(fetchone, tp=tp)
    fetchall = partial(fetchall, tp=tp)

    ui_thread_utid = get_utid_by_thread_name("awei.hmos.vmall")
    ui_track_id = fetchone(f"select id from thread_track where utid={ui_thread_utid}").id
    vsync_slices = fetchall(
        f"select ts, dur from slice where track_id = {ui_track_id} and name like 'H:OnVsyncCallback %' order by ts asc"
    )

    for prev_slice, curr_slice in tqdm(list(zip(vsync_slices[:-1], vsync_slices[1:]))):
        did_mounts = fetchall(
            f"select name from slice where track_id = {ui_track_id} and name like 'H:#RNOH::MountingManager::didMount %' and ts >= {prev_slice.ts + prev_slice.dur} and ts + dur < {curr_slice.ts} order by ts asc"
        )
        # TODO: 是不是真的jank了，当时目标帧率多少
        did_mounts = list(did_mounts)
        if len(did_mounts) == 0:
            continue

        # H:#RNOH::MountingManager::didMount 1
        mount_counts = [int(s.name.split(" ")[1]) for s in did_mounts]
        create_count = sum(c for c in mount_counts if c <= 70)
        other_count = sum(c for c in mount_counts if c > 70)

        # print(mount_counts, create_count, other_count)
        stats.append(
            {
                "create": create_count,
                "other": other_count,
                "vsync": curr_slice.dur,
            }
        )

# print(stats)
json.dump(stats, open("res.json", "w"))
